"""
进程清理脚本 - 行业标准实践

遵循Unix/Windows进程管理最佳实践：
1. 优雅关闭 (Graceful Shutdown)
2. 超时强制终止 (Timeout Force Kill)
3. 端口检查和释放
4. 资源清理验证
"""

import os
import sys
import time
import signal
import psutil
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class ProcessManager:
    """进程管理器 - 行业标准实践"""

    def __init__(self, port: int = 8000, process_name: str = "python"):
        self.port = port
        self.process_name = process_name
        self.grace_period = 5  # 优雅关闭等待时间（秒）
        self.force_timeout = 10  # 强制终止超时（秒）

    def find_processes_by_port(self) -> list[psutil.Process]:
        """通过端口查找占用进程"""
        processes = []
        for conn in psutil.net_connections(kind='inet'):
            if conn.laddr.port == self.port and conn.status == 'LISTEN':
                try:
                    process = psutil.Process(conn.pid)
                    processes.append(process)
                    logger.info(f"Found process {conn.pid} listening on port {self.port}")
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass
        return processes

    def find_processes_by_name(self) -> list[psutil.Process]:
        """通过名称查找进程"""
        processes = []
        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
            try:
                # 检查是否是Python进程且在backend目录
                if self.process_name in proc.info['name'].lower():
                    cmdline = ' '.join(proc.info['cmdline'] or [])
                    if 'uvicorn' in cmdline or 'app.main' in cmdline:
                        processes.append(proc)
                        logger.info(f"Found related process {proc.pid}: {cmdline[:100]}")
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
        return processes

    def graceful_shutdown(self, process: psutil.Process) -> bool:
        """
        优雅关闭进程

        最佳实践：
        1. 先发送 SIGTERM (15) - 允许进程清理资源
        2. 等待 grace_period 秒
        3. 如果进程仍存在，发送 SIGKILL (9)
        """
        try:
            pid = process.pid
            logger.info(f"Attempting graceful shutdown of process {pid}")

            # Step 1: 发送 SIGTERM
            if sys.platform == 'win32':
                # Windows: 使用 terminate()
                process.terminate()
            else:
                # Unix: 发送 SIGTERM
                process.send_signal(signal.SIGTERM)

            # Step 2: 等待进程优雅退出
            try:
                process.wait(timeout=self.grace_period)
                logger.info(f"Process {pid} terminated gracefully")
                return True
            except psutil.TimeoutExpired:
                logger.warning(f"Process {pid} did not terminate gracefully, forcing...")

                # Step 3: 强制终止
                process.kill()
                try:
                    process.wait(timeout=self.force_timeout)
                    logger.info(f"Process {pid} forcefully killed")
                    return True
                except psutil.TimeoutExpired:
                    logger.error(f"Failed to kill process {pid}")
                    return False

        except psutil.NoSuchProcess:
            logger.info(f"Process {pid} already terminated")
            return True
        except Exception as e:
            logger.error(f"Error shutting down process {pid}: {e}")
            return False

    def verify_port_released(self, max_retries: int = 10) -> bool:
        """验证端口是否释放（处理TIME_WAIT状态）"""
        for i in range(max_retries):
            processes = self.find_processes_by_port()
            if not processes:
                logger.info(f"Port {self.port} is now free")
                return True

            logger.info(f"Port {self.port} still occupied, waiting... ({i+1}/{max_retries})")
            time.sleep(1)

        logger.error(f"Port {self.port} still occupied after {max_retries} retries")
        return False

    def cleanup_all(self) -> dict:
        """
        完整清理流程

        返回清理结果统计
        """
        stats = {
            'port_processes': 0,
            'name_processes': 0,
            'graceful': 0,
            'forced': 0,
            'failed': 0,
            'port_released': False
        }

        # 1. 查找并关闭占用端口的进程
        logger.info(f"=== Step 1: Cleaning processes on port {self.port} ===")
        port_processes = self.find_processes_by_port()
        stats['port_processes'] = len(port_processes)

        for proc in port_processes:
            if self.graceful_shutdown(proc):
                stats['graceful'] += 1
            else:
                stats['failed'] += 1

        # 2. 查找并关闭相关Python进程
        logger.info(f"=== Step 2: Cleaning related {self.process_name} processes ===")
        name_processes = self.find_processes_by_name()
        stats['name_processes'] = len(name_processes)

        for proc in name_processes:
            # 避免重复关闭
            if proc.pid not in [p.pid for p in port_processes]:
                if self.graceful_shutdown(proc):
                    stats['graceful'] += 1
                else:
                    stats['failed'] += 1

        # 3. 验证端口释放
        logger.info("=== Step 3: Verifying port release ===")
        stats['port_released'] = self.verify_port_released()

        return stats


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description='Clean up backend processes')
    parser.add_argument('--port', type=int, default=8000, help='Port to clean (default: 8000)')
    parser.add_argument('--process-name', default='python', help='Process name to clean (default: python)')
    parser.add_argument('--dry-run', action='store_true', help='List processes without killing')

    args = parser.parse_args()

    manager = ProcessManager(port=args.port, process_name=args.process_name)

    if args.dry_run:
        logger.info("=== DRY RUN MODE ===")
        port_procs = manager.find_processes_by_port()
        name_procs = manager.find_processes_by_name()

        print(f"\nProcesses on port {args.port}:")
        for proc in port_procs:
            print(f"  PID {proc.pid}: {' '.join(proc.cmdline()[:5])}")

        print(f"\nRelated {args.process_name} processes:")
        for proc in name_procs:
            print(f"  PID {proc.pid}: {' '.join(proc.cmdline()[:5])}")
    else:
        stats = manager.cleanup_all()

        print("\n" + "=" * 60)
        print("CLEANUP SUMMARY")
        print("=" * 60)
        print(f"Port processes found:    {stats['port_processes']}")
        print(f"Name processes found:    {stats['name_processes']}")
        print(f"Gracefully terminated:   {stats['graceful']}")
        print(f"Forcefully terminated:   {stats['forced']}")
        print(f"Failed to terminate:     {stats['failed']}")
        print(f"Port released:           {'YES' if stats['port_released'] else 'NO'}")
        print("=" * 60)

        if not stats['port_released']:
            logger.error("Port is still occupied. Manual intervention may be required.")
            sys.exit(1)

        if stats['failed'] > 0:
            logger.warning(f"{stats['failed']} process(es) failed to terminate")
            sys.exit(1)

        logger.info("Cleanup completed successfully!")
        sys.exit(0)


if __name__ == '__main__':
    main()
